import argparse
import subprocess
import sys
from pathlib import Path

def run_command(cmd, **kwargs):
    result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, **kwargs)
    return result.returncode, result.stdout, result.stderr

def ensure_dir(path: Path):
    path.mkdir(parents=True, exist_ok=True)

def main():
    parser = argparse.ArgumentParser(description="Evaluate RISC-V sysy testcases.")
    parser.add_argument("testcase_dir", type=Path, help="Path to testcases.")
    parser.add_argument("mode", choices=["debug", "test"], help="Run mode.")
    parser.add_argument("--answers", type=Path, default=None, help="Path to standard answer directory (default: auto detect).")
    args = parser.parse_args()

    test_dir = args.testcase_dir.resolve()
    debug_mode = args.mode == "debug"

    if not test_dir.exists():
        print(f"[x] PATH ERROR: {test_dir} does not exist.")
        sys.exit(1)

    project_dir = Path(__file__).resolve().parent.parent.parent.parent
    io_file = project_dir / "src/io/io.c"
    compiler = "sysy"
    answer_dir = args.answers.resolve() if args.answers else (project_dir / "tests/4-code-gen/autogen/answers")
    output_dir = Path("output")
    log_file = Path("log.txt")
    log_file.write_text("")

    print(f"[info] Start testing, using testcase dir: {test_dir}")
    print(f"[info] Answer path: {answer_dir}")

    for level_dir in sorted(test_dir.glob("lv*")):
        if not level_dir.is_dir():
            continue
        level_name = level_dir.name
        level_output = output_dir / level_name
        ensure_dir(level_output)

        for case_file in sorted(level_dir.glob("*.sysy")):
            base_name = case_file.stem
            std_out_file = answer_dir / level_name / f"{base_name}.out"
            in_file = level_dir / f"{base_name}.in"
            asm_file = level_output / f"{base_name}.s"
            exe_file = level_output / base_name
            out_file = level_output / f"{base_name}.out"
            ll_file = level_output / f"{base_name}.ll"

            with log_file.open("a") as log:
                log.write(f"=========={case_file}==========\n")

            print(f"{level_name}/{base_name}...", end="")

            # LLVM IR
            if debug_mode:
                ret, _, err = run_command([str(compiler), "-emit-llvm", "-mem2reg", str(case_file), "-o", str(ll_file)])
                if ret != 0:
                    print("\033[1;31mCE\033[0m")
                    continue

            # Compile to .s
            ret, _, err = run_command([str(compiler), "-S","-mem2reg", str(case_file), "-o", str(asm_file)])
            if ret != 0:
                print("\033[1;31mCE\033[0m (sysy)")
                continue

            # Assemble
            ret, _, err = run_command(["riscv64-unknown-elf-gcc", "-static", str(asm_file), str(io_file), "-o", str(exe_file)])
            if ret != 0:
                print("\033[1;31mCE\033[0m (gcc)")
                continue

            # Run with QEMU
            with out_file.open("w") as out_f:
                try:
                    if in_file.exists():
                        with in_file.open("r") as in_f:
                            proc = subprocess.run(["qemu-riscv64", str(exe_file)],
                                                  stdin=in_f, stdout=out_f, stderr=subprocess.DEVNULL)
                    else:
                        proc = subprocess.run(["qemu-riscv64", str(exe_file)],
                                              stdout=out_f, stderr=subprocess.DEVNULL)
                    return_code = proc.returncode
                except Exception as e:
                    print(f"\033[1;31mERR\033[0m (qemu: {e})")
                    continue

			# Append return value
            with out_file.open("r") as f:
                output_content = f.read()

            with out_file.open("a") as f:
                if output_content and not output_content.endswith('\n'):
                    f.write("\n")  # Ensure separation
                f.write(str(return_code) + "\n")

            # Compare
            try:
                with open(std_out_file, "r") as f1, open(out_file, "r") as f2:
                    std_lines = [line.rstrip('\r\n') for line in f1]
                    out_lines = [line.rstrip('\r\n') for line in f2]

                    with log_file.open("a") as log:
                        log.write(f"[Output] Your Output:\n")
                        log.write("\n".join(out_lines) + "\n")
                        log.write(f"[Output] Expected Output:\n")
                        log.write("\n".join(std_lines) + "\n")

                    if std_lines != out_lines:
                        print("\033[1;31mWA\033[0m")
                        with log_file.open("a") as log:
                            log.write(f"[WA] Output mismatch: {std_out_file} vs {out_file}\n")
                        continue
            except FileNotFoundError:
                print(f"\033[1;31mNO ANSWER\033[0m")
                continue

            print("\033[1;32mOK\033[0m")

if __name__ == "__main__":
    main()